#include <assert.h>
#include <inttypes.h>
#include <stdarg.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>

typedef uint32_t u32;
typedef int32_t i32;
typedef uint16_t u16;
typedef int16_t i16;

enum log_level {
    LOG_DEBUG,
    LOG_INFO,
    LOG_ERROR,
};

static inline void logging(enum log_level level, char const* fmt, ...)
{
    va_list ap;
    va_start(ap, fmt);

    static char const* name[] = {
        [LOG_DEBUG] = "DEBUG",
        [LOG_INFO] = "INFO ",
        [LOG_ERROR] = "ERROR",
    };

    fprintf(stderr, "%s: ", name[level]);
    vfprintf(stderr, fmt, ap);
    fprintf(stderr, "\n");
    fflush(stderr);

    va_end(ap);
}

#define ARRAY_LEN(A) sizeof(A) / sizeof(*A)

struct vm {
    u32 stack[1024];
    u32 len;
};

#define VM_INIT { .len = 0 }

// Operation code of the VM.
//
// Comments before each opcode use the following notations:
//
// [a b c ... OP]
//   OP will pop arguments `a`, `b`, `c`, ... at top of stack. The last
//   argument is that at the most top. Then push a single result to the stack
//   top.
//
// [a b c ... OP(arg)]
//   OP contains a inline argument `arg` in operation code.
//
// _ = [...]
//   There are no result to push into stack.
//
// x, y, z, ... = [...]
//   There are more then one results to push into stack. The last result will
//   at the most top.
//
enum opcode {

#define OP_BINARY_XMAP(_) _(ADD, +) _(SUB, -) _(MUL, *) _(DIV, /) _(MOD, %)

#define _(NAME, OPERATOR) OP_##NAME, OP_##NAME##_KEEP,
    // OP_{NAME}: [a, b, {NAME}]
    // OP_{NAME}_KEEP: a, b, result = [a, b, {NAME}]
    OP_BINARY_XMAP(_)
#undef _

    // Push number into stack top.
    //
    //   [push(x)]
    OP_PUSH,

    // _ = [x, pop]
    OP_POP,

    // y = [x, y, rm1]
    OP_RM1,

    // y, z = [x, y, z, rm2]
    OP_RM2,

    // Print all stack.
    //
    //   _ = [pstack]
    OP_PSTACK,

    // Print stack top.
    //
    //   _ = [ptop]
    OP_PTOP,

    // [goto(label)]
    OP_GOTO,

    // If stack top is zero, jump to a label.
    //
    //   v = [v eq0(label_id)]
    OP_IF_EQ0,

    // Here is the label.
    //
    //   _ = [label(label_id)]
    OP_LABEL,

    // End of runtime.
    //
    //   _ = [end]
    OP_END = 255,
};

static inline enum opcode
op_if_eq0(u16 id)
{
    return OP_IF_EQ0 | ((u32)id << 8);
}

static inline enum opcode
op_goto(u16 id)
{
    return OP_GOTO | ((u32)id << 8);
}

static inline enum opcode
op_push(u16 n)
{
    return OP_PUSH | ((u32)n << 8);
}

static inline enum opcode
op_label(u16 id)
{
    return OP_LABEL | ((u32)id << 8);
}

enum error {
    ERR_OK,
    ERR_OUT_OF_RANGE,
};

static enum error
runtime(struct vm* vm, u32* code, u32 len)
{
#define BINARY_OP(OP)                                                          \
    vm->stack[vm->len - 2] = vm->stack[vm->len - 1] OP vm->stack[vm->len - 2]; \
    vm->len--

    u32* label = (u32*)malloc(len * sizeof(u32));
    u32 label_id;

    for (u32 i = 0; i < len; ++i) {
        switch ((enum opcode)(code[i] & 0xff)) {
        case OP_LABEL:
            label_id = code[i] >> 8;

            if (label_id >= len) {
                logging(LOG_ERROR, "label id out of range");
                return ERR_OUT_OF_RANGE;
            }

            label[label_id] = i;
            break;
        default:
            break;
        }
    }

    static char const* opname[] = {
        [OP_LABEL] = "label",
        [OP_MOD_KEEP] = "mod_keep",

        [OP_PUSH] = "push",
        [OP_POP] = "pop",
        [OP_RM1] = "rm1",
        [OP_RM2] = "rm2",

        [OP_GOTO] = "goto",
        [OP_IF_EQ0] = "if_eq0",

        [OP_PSTACK] = "pstack",
        [OP_PTOP] = "ptop",

        [OP_END] = "end",
    };

    bool stop = false;
    for (u32 i = 0; !stop; ++i) {
        enum opcode opcode = (enum opcode)(code[i] & 0xff);

        logging(LOG_DEBUG, "code: %" PRIu32 ": %s", i, opname[opcode]);

        switch (opcode) {

#define _(NAME, OPERATOR)                                                                  \
    case OP_##NAME:                                                                        \
        vm->stack[vm->len - 2] = (vm->stack[vm->len - 2] OPERATOR vm->stack[vm->len - 1]); \
        vm->len--;                                                                         \
        break;                                                                             \
    case OP_##NAME##_KEEP:                                                                 \
        vm->stack[vm->len] = (vm->stack[vm->len - 2] OPERATOR vm->stack[vm->len - 1]);     \
        vm->len++;                                                                         \
        break;

            OP_BINARY_XMAP(_)
#undef _

        case OP_PUSH:
            vm->stack[vm->len++] = (code[i] >> 8);
            break;
        case OP_POP:
            vm->len--;
            break;
        case OP_RM1:
            // y = [x, y, rm1]
            vm->stack[vm->len - 2] = vm->stack[vm->len - 1];
            vm->len--;
            break;
        case OP_RM2:
            // y, z = [x, y, z, rm2]
            vm->stack[vm->len - 3] = vm->stack[vm->len - 2];
            vm->stack[vm->len - 2] = vm->stack[vm->len - 1];
            vm->len--;
            break;

        case OP_GOTO:
            i = label[code[i] >> 8];
            break;

        case OP_IF_EQ0:
            logging(LOG_DEBUG, "  [v=%" PRIu32 ", label=%" PRIu32 "->%" PRIu32 ", eq]",
                vm->stack[vm->len - 1],
                code[i] >> 8,
                label[code[i] >> 8]);

            if (vm->stack[vm->len - 1] == 0) {
                i = label[code[i] >> 8];
            }
            break;
        case OP_PSTACK:
            logging(LOG_INFO, "pstack: [");
            for (u32 i = 0; i < vm->len; ++i) {
                logging(LOG_INFO, "pstack:   %" PRIu32 ",", vm->stack[i]);
            }
            logging(LOG_INFO, "pstack: ]");

            break;
        case OP_PTOP:
            logging(LOG_INFO, "ptop: %" PRIu32 "", vm->stack[vm->len - 1]);
            break;
        case OP_LABEL:
            break;
        case OP_END:
            stop = true;
            break;
        }
    }

    free(label);

    return ERR_OK;
}

static void
test_add(void)
{
    u32 code[] = {
        op_push(41 * 3),
        op_push(41 * 2),
        OP_ADD,
        OP_END,
    };

    struct vm vm = VM_INIT;

    runtime(&vm, code, ARRAY_LEN(code));

    assert(vm.len == 1);
    assert(vm.stack[0] == 41 * 5);
}

static void
test_gcd(void)
{
    enum label {
        LABEL_RETURN_A,
        LABEL_GCD,
    };

    u32 code[] = {
        op_push(41 * 3), // a
        op_push(41 * 2), // b

        op_label(LABEL_GCD),
        op_if_eq0(LABEL_RETURN_A),

        OP_MOD_KEEP, // [a, b, a % b]
        OP_RM2, // [b, a % b]
        op_goto(LABEL_GCD),

        op_label(LABEL_RETURN_A),
        OP_POP, // [a, b] -> a

        OP_END,
    };

    struct vm vm = VM_INIT;

    runtime(&vm, code, ARRAY_LEN(code));

    assert(vm.len == 1);
    assert(vm.stack[0] == 41);
}

static void
test_jump(void)
{
    struct vm vm = VM_INIT;

    enum label {
        LABEL_END = 10,
    };

    u32 code[1024] = {
        op_push(0),
        op_if_eq0(LABEL_END),

        op_push(1), // This will be skiped

        op_label(LABEL_END),
        OP_END,
    };

    runtime(&vm, code, ARRAY_LEN(code));
    assert(vm.len == 1);
    assert(vm.stack[0] == 0);
}

int main(void)
{
    // test_add();
    // test_jump();
    test_gcd();

    return 0;
}
